import torch
import numpy as np
import torch.nn as nn
from typing import List, Tuple, Optional, Callable, Union


embed_meta={'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0}
transf_layers = {
'embedding.word_embeddings.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.input_layernorm.weight':None,
'layers.0.input_layernorm.bias':None,
'layers.0.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.0.post_attention_layernorm.weight':None,
'layers.0.post_attention_layernorm.bias':None,
'layers.0.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.0.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.1.input_layernorm.weight':None,
'layers.1.input_layernorm.bias':None,
'layers.1.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.1.post_attention_layernorm.weight':None,
'layers.1.post_attention_layernorm.bias':None,
'layers.1.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.1.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.2.input_layernorm.weight':None,
'layers.2.input_layernorm.bias':None,
'layers.2.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.2.post_attention_layernorm.weight':None,
'layers.2.post_attention_layernorm.bias':None,
'layers.2.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.2.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.3.input_layernorm.weight':None,
'layers.3.input_layernorm.bias':None,
'layers.3.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.3.post_attention_layernorm.weight':None,
'layers.3.post_attention_layernorm.bias':None,
'layers.3.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.3.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.4.input_layernorm.weight':None,
'layers.4.input_layernorm.bias':None,
'layers.4.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.4.post_attention_layernorm.weight':None,
'layers.4.post_attention_layernorm.bias':None,
'layers.4.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.4.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.5.input_layernorm.weight':None,
'layers.5.input_layernorm.bias':None,
'layers.5.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.5.post_attention_layernorm.weight':None,
'layers.5.post_attention_layernorm.bias':None,
'layers.5.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.5.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.6.input_layernorm.weight':None,
'layers.6.input_layernorm.bias':None,
'layers.6.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.6.post_attention_layernorm.weight':None,
'layers.6.post_attention_layernorm.bias':None,
'layers.6.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.6.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.7.input_layernorm.weight':None,
'layers.7.input_layernorm.bias':None,
'layers.7.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.7.post_attention_layernorm.weight':None,
'layers.7.post_attention_layernorm.bias':None,
'layers.7.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.7.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.8.input_layernorm.weight':None,
'layers.8.input_layernorm.bias':None,
'layers.8.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.8.post_attention_layernorm.weight':None,
'layers.8.post_attention_layernorm.bias':None,
'layers.8.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.8.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.9.input_layernorm.weight':None,
'layers.9.input_layernorm.bias':None,
'layers.9.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.9.post_attention_layernorm.weight':None,
'layers.9.post_attention_layernorm.bias':None,
'layers.9.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.9.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.10.input_layernorm.weight':None,
'layers.10.input_layernorm.bias':None,
'layers.10.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.10.post_attention_layernorm.weight':None,
'layers.10.post_attention_layernorm.bias':None,
'layers.10.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.10.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.11.input_layernorm.weight':None,
'layers.11.input_layernorm.bias':None,
'layers.11.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.11.post_attention_layernorm.weight':None,
'layers.11.post_attention_layernorm.bias':None,
'layers.11.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.11.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.12.input_layernorm.weight':None,
'layers.12.input_layernorm.bias':None,
'layers.12.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.12.post_attention_layernorm.weight':None,
'layers.12.post_attention_layernorm.bias':None,
'layers.12.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.12.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.13.input_layernorm.weight':None,
'layers.13.input_layernorm.bias':None,
'layers.13.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.13.post_attention_layernorm.weight':None,
'layers.13.post_attention_layernorm.bias':None,
'layers.13.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.13.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.14.input_layernorm.weight':None,
'layers.14.input_layernorm.bias':None,
'layers.14.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.14.post_attention_layernorm.weight':None,
'layers.14.post_attention_layernorm.bias':None,
'layers.14.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.14.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.15.input_layernorm.weight':None,
'layers.15.input_layernorm.bias':None,
'layers.15.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.15.post_attention_layernorm.weight':None,
'layers.15.post_attention_layernorm.bias':None,
'layers.15.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.15.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.16.input_layernorm.weight':None,
'layers.16.input_layernorm.bias':None,
'layers.16.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.16.post_attention_layernorm.weight':None,
'layers.16.post_attention_layernorm.bias':None,
'layers.16.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.16.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.17.input_layernorm.weight':None,
'layers.17.input_layernorm.bias':None,
'layers.17.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.17.post_attention_layernorm.weight':None,
'layers.17.post_attention_layernorm.bias':None,
'layers.17.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.17.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.18.input_layernorm.weight':None,
'layers.18.input_layernorm.bias':None,
'layers.18.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.18.post_attention_layernorm.weight':None,
'layers.18.post_attention_layernorm.bias':None,
'layers.18.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.18.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.19.input_layernorm.weight':None,
'layers.19.input_layernorm.bias':None,
'layers.19.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.19.post_attention_layernorm.weight':None,
'layers.19.post_attention_layernorm.bias':None,
'layers.19.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.19.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.20.input_layernorm.weight':None,
'layers.20.input_layernorm.bias':None,
'layers.20.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.20.post_attention_layernorm.weight':None,
'layers.20.post_attention_layernorm.bias':None,
'layers.20.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.20.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.21.input_layernorm.weight':None,
'layers.21.input_layernorm.bias':None,
'layers.21.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.21.post_attention_layernorm.weight':None,
'layers.21.post_attention_layernorm.bias':None,
'layers.21.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.21.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.22.input_layernorm.weight':None,
'layers.22.input_layernorm.bias':None,
'layers.22.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.22.post_attention_layernorm.weight':None,
'layers.22.post_attention_layernorm.bias':None,
'layers.22.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.22.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.23.input_layernorm.weight':None,
'layers.23.input_layernorm.bias':None,
'layers.23.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.23.post_attention_layernorm.weight':None,
'layers.23.post_attention_layernorm.bias':None,
'layers.23.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.23.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.24.input_layernorm.weight':None,
'layers.24.input_layernorm.bias':None,
'layers.24.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.24.post_attention_layernorm.weight':None,
'layers.24.post_attention_layernorm.bias':None,
'layers.24.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.24.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.25.input_layernorm.weight':None,
'layers.25.input_layernorm.bias':None,
'layers.25.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.25.post_attention_layernorm.weight':None,
'layers.25.post_attention_layernorm.bias':None,
'layers.25.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.25.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.26.input_layernorm.weight':None,
'layers.26.input_layernorm.bias':None,
'layers.26.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.26.post_attention_layernorm.weight':None,
'layers.26.post_attention_layernorm.bias':None,
'layers.26.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.26.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.27.input_layernorm.weight':None,
'layers.27.input_layernorm.bias':None,
'layers.27.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.27.post_attention_layernorm.weight':None,
'layers.27.post_attention_layernorm.bias':None,
'layers.27.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.27.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.28.input_layernorm.weight':None,
'layers.28.input_layernorm.bias':None,
'layers.28.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.28.post_attention_layernorm.weight':None,
'layers.28.post_attention_layernorm.bias':None,
'layers.28.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.28.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.29.input_layernorm.weight':None,
'layers.29.input_layernorm.bias':None,
'layers.29.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.29.post_attention_layernorm.weight':None,
'layers.29.post_attention_layernorm.bias':None,
'layers.29.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.29.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.30.input_layernorm.weight':None,
'layers.30.input_layernorm.bias':None,
'layers.30.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.30.post_attention_layernorm.weight':None,
'layers.30.post_attention_layernorm.bias':None,
'layers.30.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.30.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.31.input_layernorm.weight':None,
'layers.31.input_layernorm.bias':None,
'layers.31.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.31.post_attention_layernorm.weight':None,
'layers.31.post_attention_layernorm.bias':None,
'layers.31.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.31.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.32.input_layernorm.weight':None,
'layers.32.input_layernorm.bias':None,
'layers.32.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.32.post_attention_layernorm.weight':None,
'layers.32.post_attention_layernorm.bias':None,
'layers.32.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.32.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.33.input_layernorm.weight':None,
'layers.33.input_layernorm.bias':None,
'layers.33.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.33.post_attention_layernorm.weight':None,
'layers.33.post_attention_layernorm.bias':None,
'layers.33.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.33.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.34.input_layernorm.weight':None,
'layers.34.input_layernorm.bias':None,
'layers.34.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.34.post_attention_layernorm.weight':None,
'layers.34.post_attention_layernorm.bias':None,
'layers.34.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.34.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.35.input_layernorm.weight':None,
'layers.35.input_layernorm.bias':None,
'layers.35.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.35.post_attention_layernorm.weight':None,
'layers.35.post_attention_layernorm.bias':None,
'layers.35.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.35.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.36.input_layernorm.weight':None,
'layers.36.input_layernorm.bias':None,
'layers.36.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.36.post_attention_layernorm.weight':None,
'layers.36.post_attention_layernorm.bias':None,
'layers.36.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.36.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.37.input_layernorm.weight':None,
'layers.37.input_layernorm.bias':None,
'layers.37.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.37.post_attention_layernorm.weight':None,
'layers.37.post_attention_layernorm.bias':None,
'layers.37.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.37.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.38.input_layernorm.weight':None,
'layers.38.input_layernorm.bias':None,
'layers.38.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.38.post_attention_layernorm.weight':None,
'layers.38.post_attention_layernorm.bias':None,
'layers.38.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.38.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.39.input_layernorm.weight':None,
'layers.39.input_layernorm.bias':None,
'layers.39.self_attention.query_key_value.weight':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.self_attention.query_key_value.bias':{'partition_dim': 0, 'partition_stride': 3, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.self_attention.dense.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.self_attention.dense.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'layers.39.post_attention_layernorm.weight':None,
'layers.39.post_attention_layernorm.bias':None,
'layers.39.mlp.dense_h_to_4h.weight':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.mlp.dense_h_to_4h.bias':{'partition_dim': 0, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.mlp.dense_4h_to_h.weight':{'partition_dim': 1, 'partition_stride': 1, 'parallel_size': 4, 'parallel_index': 0},
'layers.39.mlp.dense_4h_to_h.bias':{'partition_dim': -1, 'partition_stride': 1, 'parallel_size': -1, 'parallel_index': -1},
'final_layernorm.weight':None,
'final_layernorm.bias':None
}

def ensure_divisibility(numerator: int, denominator: int) -> None:
    """Ensure that numerator is divisible by the denominator."""
    assert numerator % denominator == 0, \
        f'{numerator} is not divisible by {denominator}'

def divide(numerator: int, denominator: int) -> int:
    """Ensure that numerator is divisible by the denominator and return
    the division value."""
    ensure_divisibility(numerator, denominator)
    return numerator // denominator

def get_1d_parallel_split_size(dim_size,
                               partition_stride,
                               world_size):
    """If we split tensor by world size and stride along a dim,
    this function will return the number of chunks and size of each chunk."""

    per_partition_size = divide(dim_size, world_size)
    per_partition_per_stride_size = divide(per_partition_size, partition_stride)
    num_chunks = partition_stride * world_size
    return num_chunks, per_partition_per_stride_size
def get_1d_parallel_split_indices(rank, 
                                  partition_stride,
                                  world_size):
    """If we split tensor by world size and stride along a dim,
    this function will return the indices of splited chunks on
    the current member in TP group."""

    num_chunks = partition_stride * world_size
    return tuple(range(rank, num_chunks, world_size))

def split_tensor_to_local(tensor: torch.Tensor,
                          partition_dim:int,
                          partition_stride: Optional[int],
                          out:Optional[torch.Tensor]=None,
                          world_size: Optional[int]=None,
                          cur_rank: Optional[int]=None) -> torch.Tensor:

    """Split a tensor by tensor parallel size and stride along a dim.

    If we want to split a global tensor along `dim` to each member of tensor parallel group,
    we first split weight to `world_size*stride` partitions, each partition size equal to 
    `dim_size//(world_size*stride)`.

    with world_size=4, stride=1, each member of tensor parallel groups gets the indices of partitions:
        rank=0: [0]
        rank=1: [1]
        rank=2: [2]
        rank=3: [3]
    with world_size=4, stride=2, each member of tensor parallel groups gets the indices of partitions:
        rank=0: [0, 4]
        rank=1: [1, 5]
        rank=2: [2, 6]
        rank=3: [3, 7]
    with world_size=4, stride=4, each member of tensor parallel groups gets the indices of partitions:
        parallel rank=0: [0, 4, 8, 12]
        parallel rank=1: [1, 5, 9, 13]
        parallel rank=2: [2, 6, 10, 14]
        parallel rank=3: [3, 7, 11, 15]
    """

    if world_size == 1:
        return tensor if out is None else out.data.copy_(tensor)
    num_chunks, per_partition_per_stride_size = \
        get_1d_parallel_split_size(dim_size=tensor.size(partition_dim),
                                   partition_stride=partition_stride,
                                   world_size=world_size)
    assert per_partition_per_stride_size * num_chunks == tensor.size(partition_dim)
    tensor_chunks = torch.chunk(tensor.data, num_chunks, dim=partition_dim)
    partition_indices = get_1d_parallel_split_indices(rank=cur_rank,
                                                      partition_stride=partition_stride,
                                                      world_size=world_size)
    if out is not None and isinstance(out, torch.Tensor):
        return torch.cat([tensor_chunks[i] for i in partition_indices], dim=partition_dim, out=out)
    tensor_partition = torch.cat([tensor_chunks[i] for i in partition_indices], dim=partition_dim)
    return tensor_partition.contiguous()


def gather_tensors_to_global(tensors: List[torch.Tensor],
                             partition_dim: int,
                             partition_stride: Optional[int]=1,
                             world_size: Optional[int]=None) -> torch.Tensor:
    """Gather local tensors to a global tensor(completed), the tensors must ordered by ranks"""

    if world_size == 1:
        return tensors[0]
    local_tensor = tensors[0]
    assert len(tensors) > 0 and len(tensors) == world_size
    global_dim_size = sum([t.size(partition_dim) for t in tensors])
    assert world_size * local_tensor.size(partition_dim) == global_dim_size
    global_shape = list(local_tensor.size())
    global_shape[partition_dim] = global_dim_size
    global_tensor = torch.zeros(global_shape,
                                dtype=local_tensor.dtype,
                                device=local_tensor.device)
    num_chunks, per_size = get_1d_parallel_split_size(global_dim_size, partition_stride, world_size)
    chunks = torch.chunk(global_tensor, num_chunks, dim=partition_dim)
    partition_chunks = [torch.chunk(t, partition_stride, dim=partition_dim) for t in tensors]
    for rank in range(world_size):
        for i, index in enumerate(get_1d_parallel_split_indices(rank, partition_stride, world_size)):
            chunks[index].data.copy_(partition_chunks[rank][i].data)
    return global_tensor